// --------------------------------------------------------------------------------------------------------------------
// <copyright file="LinearSystem.cs" company="Microsoft Corporation">
// The MIT License (MIT)
// 
// Copyright (c) 2014, Microsoft Corporation
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in
//  all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
//  THE SOFTWARE.
// </copyright>
// --------------------------------------------------------------------------------------------------------------------
namespace Microsoft.Robotics.Numerics.GraphOptimization
{
    using System;
    using System.Collections.Generic;
    using Microsoft.Robotics.Numerics;
    using ShoNS.Array;

    /// <summary>
    /// Creates the linear system of a graph
    /// </summary>
    /// <typeparam name="NodeType">Type of nodes</typeparam>
    /// <typeparam name="EdgeType">Type of edges</typeparam>
    public class LinearSystem<NodeType, EdgeType> where EdgeType : Edge<NodeType>
    {
        /// <summary>
        /// A matrix
        /// </summary>
        private double[,] matrixA;

        /// <summary>
        /// B matrix
        /// </summary>
        private double[,] matrixB;

        /// <summary>
        /// Dimension of the error function
        /// </summary>
        private int dim;

        /// <summary>
        /// Initializes a new instance of the LinearSystem class.  Currently only can create a linear system for PoseSLAM
        /// Creates H and B in the linear system $Hx = B$.
        /// See $\\robofile\Baker\Shared\Users\bpeasley\g2o.pdf$ for description of this approach
        /// web link: $http://ais.informatik.uni-freiburg.de/publications/papers/kuemmerle11icra.pdf $
        /// </summary>
        /// <param name="graph">Graph to create linear system from</param>
        /// <param name="error">Pointer to error function</param>
        /// <param name="increment">Pointer to increment function</param>
        public LinearSystem(Graph<NodeType, EdgeType> graph, ErrorFunction error, IncrementFunction increment)
        {          
            this.dim = graph.Edges[0].Dimensions;
            this.Size = this.dim * (graph.Nodes.Count - 1);
            
            this.matrixA = new double[this.dim, this.dim];
            this.matrixB = new double[this.dim, this.dim];
            this.MatrixH = new double[this.Size * this.Size];
            this.VectorB = new double[this.Size];
            
            double[] edgeError = new double[this.dim];
            double[] subVectorB = new double[this.dim];

            // Loop through each edge in the graph
            for (int i = 0; i < graph.Edges.Count; i++)
            {
                // Numerically compute the Jacobian for an edge by finite differencing
                this.ComputeJacobianIJ(graph.Nodes[graph.Edges[i].IndexI], graph.Nodes[graph.Edges[i].IndexJ], graph.Edges[i], error, increment);

                // Compute the error of the edge given the current state
                error(graph.Nodes[graph.Edges[i].IndexI], graph.Nodes[graph.Edges[i].IndexJ], graph.Edges[i], ref edgeError);

                int indexI = (graph.Edges[i].IndexI - 1) * this.dim;
                int indexJ = (graph.Edges[i].IndexJ - 1) * this.dim;

                // Update the information matrix
                // See paper page 4 bottom of left column ( Hij = .... )
                // Covariance matrices are currently ignored, on todo list
                if (indexI >= 0 && indexJ >= 0)
                {
                    // Neither I nor J are the $0th$ node, which is fixed
                    this.AddSubMatrix(this.matrixA.Transpose().Multiply(this.matrixA), indexI, indexI);
                    this.AddSubMatrix(this.matrixB.Transpose().Multiply(this.matrixB), indexJ, indexJ);
                    this.AddSubMatrix(this.matrixA.Transpose().Multiply(this.matrixB), indexI, indexJ);
                    this.AddSubMatrix(this.matrixB.Transpose().Multiply(this.matrixA), indexJ, indexI);

                    subVectorB = this.matrixA.Transpose().Multiply(edgeError);
                    for (int j = indexI, row = 0; j <= indexI + 2; j++, row++)
                    {
                        this.VectorB[j] += subVectorB[row];
                    }

                    subVectorB = this.matrixB.Transpose().Multiply(edgeError);
                    for (int j = indexJ, row = 0; j <= indexJ + 2; j++, row++)
                    {
                        this.VectorB[j] += subVectorB[row];
                    }
                }
                else if (indexJ < 0) 
                {
                    // Index J is the $0th$ node, which is fixed
                    this.AddSubMatrix(this.matrixA.Transpose().Multiply(this.matrixA), indexI, indexI);
                    subVectorB = this.matrixA.Transpose().Multiply(edgeError);
                    for (int j = indexI, row = 0; j <= indexI + 2; j++, row++)
                    {
                        this.VectorB[j] += subVectorB[row];
                    }
                }
                else if (indexI < 0) 
                {
                    // Index I is the $0th$ node, which is fixed
                    this.AddSubMatrix(this.matrixB.Transpose().Multiply(this.matrixB), indexJ, indexJ);
                    subVectorB = this.matrixB.Transpose().Multiply(edgeError);
                    for (int j = indexJ, row = 0; j <= indexJ + 2; j++, row++)
                    {
                        this.VectorB[j] += subVectorB[row];
                    }
                }
            }
        }

        /// <summary>
        /// Delegate of function that computes the error
        /// </summary>
        /// <param name="state1">State one</param>
        /// <param name="state2">State two</param>
        /// <param name="measurement">Measurement between the two states</param>
        /// <param name="error">Reference to the error</param>
        public delegate void ErrorFunction(NodeType state1, NodeType state2, EdgeType measurement, ref double[] error);

        /// <summary>
        /// Delegate of function that increments the state
        /// </summary>
        /// <param name="state">state to increment</param>
        /// <param name="increment">amount to increment</param>
        /// <returns>incremented state</returns>
        public delegate NodeType IncrementFunction(NodeType state, double[] increment);

        /// <summary>
        /// Gets the H matrix
        /// </summary>
        public double[] MatrixH { get; private set; }

        /// <summary>
        /// Gets the B vector
        /// </summary>
        public double[] VectorB { get; private set; }

        /// <summary>
        /// Gets the length of a row/col of H
        /// </summary>
        public int Size { get; private set; }

        /// <summary>
        /// Solves the Linear System using $Cholesky$ Decomposition
        /// </summary>
        /// <returns>The solution to the standard linear system</returns>
        public DoubleArray Solve()
        {
            SparseDoubleArray h = this.GetH();
            DoubleArray b = this.GetB();

            // If this gives a PARDISO error of -4, it probably means that the linear system isn't positive semi definite
            SparseCholesky cholesky = new SparseCholesky(h);
            return cholesky.Solve(b.Transpose());
        }

        /// <summary>
        /// Adds a square sub matrix to $MatrixH$ with the sub matrix top left corner at index (top, left)
        /// </summary>
        /// <param name="matrix">the sub matrix to add</param>
        /// <param name="top">top corner</param>
        /// <param name="left">left corner</param>
        private void AddSubMatrix(double[,] matrix, int top, int left)
        {
            int dimSize = matrix.GetLength(0);
            for (int i = top, row = 0; i < top + dimSize; i++, row++)
            {
                for (int j = left, col = 0; j < left + dimSize; j++, col++)
                {
                    int index = i * this.Size + j;
                    this.MatrixH[i * this.Size + j] += matrix[row, col];
                }
            }
        }

        /// <summary>
        /// Numerically computes the Jacobian of connecting nodes $I$ and $J$ by finite differencing
        /// Currently only works for systems that have decoupled state spaces (i.e. Global State Space)
        /// Make small perturbations to the states of nodes $I$ and $J$ and compute errors to find $Jij$
        /// See paper page 4 middle of left column ($Jij$ = ...)
        /// </summary>
        /// <param name="state1">State one</param>
        /// <param name="state2">State two</param>
        /// <param name="measurement">measurement between the two states</param>
        /// <param name="error">Pointer to error function</param>
        /// <param name="increment">Pointer to increment function</param>
        private void ComputeJacobianIJ(NodeType state1, NodeType state2, EdgeType measurement, ErrorFunction error, IncrementFunction increment)
        {
            double[] edgeError1 = new double[this.dim];
            double[] edgeError2 = new double[this.dim];

            double step = 0.01;

            // Loop through all dimensions in state1 (I), peturb that state by the same amount (step) positively and negatively
            // Compute the error during the perturbations.  Subtract the 2 errors to get the derivative of the error
            for (int i = 0; i < measurement.Dimensions; i++)
            {
                NodeType new_state = increment(state1, measurement.ScaledBasisVector(i, -step));
                error(new_state, state2, measurement, ref edgeError1);
                new_state = increment(state1, measurement.ScaledBasisVector(i, step));
                error(new_state, state2, measurement, ref edgeError2);

                for (int j = 0; j < measurement.Dimensions; j++)
                {
                    this.matrixA[j, i] = (edgeError2[j] - edgeError1[j]) / (2.0 * step);
                }
            }

            // Same thing but for state2 (J)
            for (int i = 0; i < measurement.Dimensions; i++)
            {
                NodeType new_state = increment(state2, measurement.ScaledBasisVector(i, -step));
                error(state1, new_state, measurement, ref edgeError1);
                new_state = increment(state2, measurement.ScaledBasisVector(i, step));
                error(state1, new_state, measurement, ref edgeError2);

                for (int j = 0; j < measurement.Dimensions; j++)
                {
                    this.matrixB[j, i] = (edgeError2[j] - edgeError1[j]) / (2.0 * step);
                }
            }
        }

            /// <summary>
        /// Creates the Sparse Double Array of the information matrix needed by SHO to do $Cholesky$
        /// </summary>
        /// <returns>Sparse double array of in the information matrix</returns>
        private SparseDoubleArray GetH()
        {
            List<int> rowIndices = new List<int>();
            List<int> colIndices = new List<int>();
            List<double> values = new List<double>();

            for (int i = 0, index = 0; i < this.Size; i++)
            {
                for (int j = 0; j < this.Size; j++, index++)
                {
                    if (this.MatrixH[index] != 0)
                    {
                        rowIndices.Add(i);
                        colIndices.Add(j);
                        values.Add(this.MatrixH[index]);
                    }
                }
            }

            return new SparseDoubleArray(this.Size, this.Size, rowIndices, colIndices, values, true);
        }

        /// <summary>
        /// Creates the Double Array of B needed by SHO to do $Cholesky$
        /// </summary>
        /// <returns>double array</returns>
        private DoubleArray GetB()
        {
            DoubleArray b = new DoubleArray(this.Size);
            for (int i = 0; i < this.Size; i++)
            {
                b[i] = this.VectorB[i];
            }

            return b;
        }
    }
}
